# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import functools
import logging
import sys

import pandas

from coin.base.datetime_util import to_timestamp_int

from coin.base.book.types import TradeSide
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.exchange.bitfinex_v2.kr_rest.product import BitfinexProduct
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.gdax.kr_rest.product import GdaxProduct

from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.dumper_base import run_from_archive


class FeatureContainer(object):
  def __init__(self, exch_symbols):
    self._exch_symbols = exch_symbols
    self._ready_exch_symbols = set()
    self._ready = False
    self._fields = {}
    self._exch_cumtrade = {}

  def fillup_prefix(self, prefix, fields, timestamp):
    fields = {(prefix + "_" + key): value for key, value in fields.items()}
    fields['timestamp'] = timestamp
    return fields

  def get_raw_fields(self, prefix, book):
    if book.exchange.startswith('Okex'):
      # step_price = 0.01 # for okexfutures
      # step_price = 0.1
      step_price = 0.2
    else:
      step_price = book.ask0().price * 5e-5  # 0.5bp
    prices = {}
    n = 10
    bnames = ["bid%d" % (n - i - 1) for i in range(n)] + \
        ["ask%d" % i for i in range(n)]
    bnamemap = {}
    for i, bname in enumerate(bnames):
      bnamemap[bname] = "%s_%s" % (i, bname)

    fields = {}
    for side, method in [("bid", book.iter_bid_book), ("ask", book.iter_ask_book)]:
      for i, (fromp, top, qty) in enumerate(method(num_level=n, step_price=step_price)):
        iname = "%s%d" % (side, i)
        name = bnamemap[iname]
        # price = 0.5 * (fromp + top)
        price = fromp
        prices[name] = {'price': fromp, 'qty': qty}
        fields[iname + '_price'] = fromp
        fields[iname + '_qty'] = qty
    fields = self.fillup_prefix(prefix, fields, book.timestamp)
    return fields

  def get_raw_trade_fields(self, prefix, trade):
    if prefix not in self._exch_cumtrade:
      self._exch_cumtrade[prefix] = 0

    assert trade.side in [TradeSide.TRADE_BUY_SIDE, TradeSide.TRADE_SELL_SIDE]
    signed_qty = trade.qty * (1 if trade.side == TradeSide.TRADE_BUY_SIDE else -1)

    self._exch_cumtrade[prefix] += signed_qty
    cum_trade_qty = self._exch_cumtrade[prefix]

    fields = {'trade_price': trade.price, 'trade_qty': signed_qty, 'cum_trade_qty': cum_trade_qty}
    fields = self.fillup_prefix(prefix, fields, trade.timestamp)
    return fields

  def get_prefix(self, exchange, symbol):
    return "_".join((exchange, symbol))

  def update_book(self, book):
    exch_symbol = (book.exchange, book.symbol)
    prefix = self.get_prefix(book.exchange, book.symbol)

    if exch_symbol in self._exch_symbols:
      self._fields.update(self.get_raw_fields(prefix, book))
      if exch_symbol not in self._ready_exch_symbols:
        self._ready_exch_symbols.add(exch_symbol)
        print(set(self._exch_symbols) - set(self._ready_exch_symbols))
      if len(self._ready_exch_symbols) == len(self._exch_symbols) and not self._ready:
        print("ready")
        self._ready = True
    else:
      raise ValueError(exch_symbol, self._exch_symbols)

  def update_trade(self, exchange, trade):
    prefix = self.get_prefix(exchange, trade.symbol)
    self._fields.update(self.get_raw_trade_fields(prefix, trade))


class TimedSampler(object):
  def __init__(self, sampling_period=None):
    self._last_sampled = 0
    self._sampling_period = sampling_period or (1 * 1e9)

  def try_sample(self, timestamp):
    if self._last_sampled + self._sampling_period < timestamp:
      self._last_sampled = timestamp
      return True
    else:
      return False

  def update_book(self, book):
    return self.try_sample(book.timestamp)

  def update_trade(self, exchange, trade):
    return self.try_sample(trade.timestamp)


class OnBookSampler(object):
  def __init__(self, product):
    self._product = product

  def update_book(self, book):
    return (book.exchange, book.symbol) == self._product

  def update_trade(self, exchange, trade):
    return False


class ExploreBook(object):
  def __init__(self, machine, products, dump_dir='.', postfix=None, start_ts=None):
    self._machine = machine
    self._products = products
    self._tick_sizes = (0.5, 0.01)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._fields_out = []
    self._fields_view = []
    exch_symbols = [('Okex', product.native_subscription_symbol)
                    if product.exchange == 'Okex' and product.product_type == 'Futures' else
                    (product.exchange, product.native_symbol) for product in products]
    self._feature_container = FeatureContainer(exch_symbols)

    # self._sampler = TimedSampler()
    self._sampler = OnBookSampler(('Okex', 'btc_quarter'))

    # self._watch_product = ('Bitmex', 'XBTUSD')
    self._watch_product = ('Okex', 'btc_quarter')
    self._time_set = set()
    print(exch_symbols)

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      subscription_key = get_subscription_key(product)
      if book_builder_name == subscription_key:
        book_builder.subscribe(product, self.on_book)
        if hasattr(book_builder, "subscribe_trade"):
          book_builder.subscribe_trade(product, functools.partial(self.on_trade, product.exchange))

  def on_book(self, book):
    if book.timestamp < self._start_ts:
      return
    intervalstamp = int(book.timestamp / 600e9) * 600e9
    if intervalstamp not in self._time_set:
      self._time_set.add(intervalstamp)
      print(pandas.DatetimeIndex([self._start_ts]))
      print(pandas.DatetimeIndex([intervalstamp])[0])

    self._feature_container.update_book(book)
    if self._feature_container._ready and self._sampler.update_book(book):
      self.on_sample()
    self.try_view(book.exchange, book.symbol)

  def on_sample(self):
    self._fields_out.append(self._feature_container._fields.copy())
    if len(self._fields_out) % 200 == 0:
      print(len(self._fields_out))
      df = pandas.DataFrame(self._fields_out)
      print(pandas.concat([df[:2], df[-2:]], axis=0))

    if len(self._fields_out) % 2000 == 0:
      pandas.DataFrame(self._fields_out).to_csv("%s.csv" % self._machine)

    # if len(self._fields_out) > 20000:
    # pandas.DataFrame(self._fields_out).to_csv("%s.csv" % self._machine)
    # end

  def on_trade(self, exchange, trade):
    if trade.timestamp < self._start_ts:
      return
    self._feature_container.update_trade(exchange, trade)
    if self._feature_container._ready and self._sampler.update_trade(exchange, trade):
      self.on_sample()
    self.try_view(exchange, trade.symbol)
    # if exchange == self._products[0].exchange:
    #  self.try_sample(trade.timestamp)

  def try_view(self, exchange, symbol):
    return
    if (exchange, symbol) == self._watch_product:
      import matplotlib.pyplot as plt
      from pylab import rcParams
      rcParams['figure.figsize'] = 16, 7
      self._fields_view.append(self._feature_container._fields.copy())
      if (len(self._fields_view) % 50 != 0):
        return

      plotdf = pandas.DataFrame(self._fields_view[-100:])

      xtime = pandas.DatetimeIndex(plotdf['timestamp'])

      prod = '_'.join(self._watch_product)
      sample_qty_map = {
          'Okex_btc_quarter': 100,
          'Bitmex_XBTUSD': 100000,
      }

      sample_qty = sample_qty_map[prod]

      for i in range(10):
        plt.plot(xtime,
                 plotdf['%s_ask%s_price' % (prod, i)],
                 'g-',
                 xtime,
                 plotdf['%s_bid%s_price' % (prod, i)],
                 'k-',
                 linewidth=0.1)
        for markersize, qty, marker in [(2.0, sample_qty * 5, 'k.'), (1.0, sample_qty, 'g.'),
                                        (0.5, sample_qty * 0.1, 'g.')]:
          idx = plotdf['%s_ask%s_qty' % (prod, i)] > qty
          xidx = plotdf.index[idx]
          xidx = xtime[xidx]
          if len(xidx) > 0:
            plt.plot(xidx,
                     plotdf.loc[idx, '%s_ask%s_price' % (prod, i)],
                     marker,
                     markersize=markersize)
          idx = plotdf['%s_bid%s_qty' % (prod, i)] > qty

          xidx = plotdf.index[idx]
          xidx = xtime[xidx]
          if len(xidx) > 0:
            plt.plot(xidx,
                     plotdf.loc[idx, '%s_bid%s_price' % (prod, i)],
                     marker,
                     markersize=markersize)
      # import pdb; pdb.set_trace()
      if '%s_trade_price' % prod in plotdf.columns:
        plt.plot(xtime, plotdf['%s_trade_price' % prod], 'r-', linewidth=0.5)
        # plt.grid(True)
        plt.twinx()
        plt.plot(xtime, plotdf['%s_cum_trade_qty' % prod], 'g.-', linewidth=0.5, markersize=0.5)
      plt.gcf().autofmt_xdate()
      plt.title(prod)
      # plt.subplot(212); plt.plot(sorted(plotdf['timestamp'].diff()))
      plt.waitforbuttonpress(timeout=-1)
      plt.clf()


def get_products(machine):
  okex_futures_product1 = OkexFuturesProduct.FromStr("BTC-USD.THIS_WEEK")
  okex_futures_product2 = OkexFuturesProduct.FromStr("BTC-USD.NEXT_WEEK")
  okex_futures_product3 = OkexFuturesProduct.FromStr("BTC-USD.QUARTER")

  bitmex_product = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')
  bitflyer_product = BitflyerFuturesProduct.FromStr("BTC-JPY.PERPETUAL")
  binance_product = BinanceProduct.FromStr("BTC-USDT")
  bitfinex_product = BitfinexProduct.FromStr("BTC-USDT")
  okex_product = OkexProduct.FromStr("BTC-USDT")
  gdax_product = GdaxProduct.FromStr("BTC-USD")

  # products = [bitmex_product, okex_futures_product1, okex_futures_product2, okex_futures_product3, binance_product, bitfinex_product, okex_product, gdax_product]

  # products = [okex_product]
  if machine == 'feed-01.ap-northeast-1.aws':
    products = [
        bitflyer_product,
        bitmex_product,
        okex_futures_product1,
        okex_futures_product2,
        okex_futures_product3,
        okex_product,  # bitfinex_product,
    ]
  elif machine == 'feed-01.eu-west-1.aws':
    products = [
        bitmex_product,
        okex_futures_product1,
        okex_futures_product2,
        okex_futures_product3,
        okex_product,
        bitfinex_product
    ]
  else:
    raise ValueError(machine)

  return products


def launch(date):
  # machine='feed-01.eu-west-1.aws',
  # machine = 'feed-01.ap-northeast-1.aws'
  for machine in [  #'feed-01.eu-west-1.aws',
      'feed-01.ap-northeast-1.aws'
  ]:
    products = get_products(machine)

    start_time = datetime.datetime(2018, 3, date, hour=16)
    hours = 24

    strategy = ExploreBook(machine=machine,
                           products=products,
                           dump_dir='dump',
                           postfix='_%s_%dH' % (start_time.strftime('%Y%m%d'), hours),
                           start_ts=start_time + datetime.timedelta(minutes=60))

    run_from_archive(products,
                     strategy.on_book_reset,
                     start_time,
                     start_time + datetime.timedelta(hours=hours),
                     machine=[machine])
  return 0


def main(argv):
  date = 17
  launch(date)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  try:
    sys.exit(main(sys.argv))
  except KeyboardInterrupt:
    sys.exit(1)
